// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/lib/kll_sketch.h>

namespace kumo::pollux::functions::kll {
    uint32_t kFromEpsilon(double eps) {
        return ceil(exp(1.0285 * log(2.296 / eps)));
    }

    namespace detail {
        namespace {
            constexpr uint8_t kMinBufferWidth = 8;

            double powerOfTwoThirds(int n) {
                static const auto kMemo = [] {
                    std::array<double, kMaxLevel> memo;
                    for (int i = 0; i < kMaxLevel; ++i) {
                        memo[i] = pow(2.0 / 3.0, i);
                    }
                    return memo;
                }();
                return kMemo[n];
            }
        } // namespace

        uint32_t computeTotalCapacity(uint32_t k, uint8_t numLevels) {
            uint32_t total = 0;
            for (uint8_t h = 0; h < numLevels; ++h) {
                total += levelCapacity(k, numLevels, h);
            }
            return total;
        }

        uint32_t levelCapacity(uint32_t k, uint8_t numLevels, uint8_t height) {
            POLLUX_DCHECK_LT(height, numLevels);
            POLLUX_DCHECK_LE(numLevels, kMaxLevel);
            return std::max<uint32_t>(
                kMinBufferWidth, k * powerOfTwoThirds(numLevels - height - 1));
        }

        uint8_t floorLog2(uint64_t p, uint64_t q) {
            for (uint8_t ans = 0;; ++ans) {
                q <<= 1;
                if (p < q) {
                    return ans;
                }
            }
        }

        uint64_t sumSampleWeights(uint8_t numLevels, const uint32_t *levels) {
            uint64_t total = 0;
            uint64_t weight = 1;
            for (uint8_t lvl = 0; lvl < numLevels; lvl++) {
                total += weight * (levels[lvl + 1] - levels[lvl]);
                weight *= 2;
            }
            return total;
        }
    } // namespace detail
} // namespace kumo::pollux::functions::kll
